package com.examples.es;

import org.apache.flink.configuration.Configuration;
import org.apache.flink.streaming.api.functions.source.RichSourceFunction;
import org.apache.http.HttpHost;
import org.elasticsearch.client.RequestOptions;
import org.elasticsearch.client.RestClient;
import org.elasticsearch.client.RestHighLevelClient;
import org.elasticsearch.client.core.CountRequest;
import org.elasticsearch.client.core.CountResponse;

import java.io.IOException;

public class IndicesCountSource extends RichSourceFunction<Integer> {
    private static final long serialVersionUID = 3334654984548091675L;


    private final String indices;
    private static RestHighLevelClient restHighLevelClient;
    private final HttpHost httpHost;
    public IndicesCountSource(HttpHost httpHost, String indices) {
        this.indices = indices;
        this.httpHost = httpHost;
    }

    @Override
    public void open(Configuration parameters) throws Exception {
        super.open(parameters);
        restHighLevelClient = new RestHighLevelClient(RestClient.builder(httpHost));
    }

    @Override
    public void run(SourceContext<Integer> sourceContext) throws Exception {
        int count = 0;
        CountRequest countRequest = new CountRequest(indices);
        try {
            CountResponse countResponse = restHighLevelClient.count(countRequest, RequestOptions.DEFAULT);
            count = Integer.parseInt(String.valueOf(countResponse.getCount()));
        } catch (IOException ignored) {}
        sourceContext.collect(count);
    }

    @Override
    public void cancel() {
        if(restHighLevelClient != null){
            try {
                restHighLevelClient.close();
            } catch (IOException e) {
                e.printStackTrace();
            }
        }
    }
}
